import cv2
import numpy as np


# 定义感兴趣区域（ROI）
def region_of_interest(image):
    height = image.shape[0]
    width = image.shape[1]

    # 定义一个区域
    mask = np.zeros_like(image)
    polygon = np.array([[
        (0, height),
        (0, height *2/3),
        (width *1/3, height /2),
        (width *2/3, height /2),
        (width, height *2/3),
        (width, height)
    ]], np.int32)

    cv2.fillPoly(mask, polygon, (255,255,255))
    masked_image = cv2.bitwise_and(image, mask)
    return masked_image

# 计算直线的斜率
def calculate_slope(line):
    x_1, y_1, x_2, y_2 = line[0]
    if x_2 - x_1 != 0:
        return (y_2 - y_1) / (x_2 - x_1)
    else:
        return np.inf

# 应用霍夫变换标记处直线
def detect_roadline(masked_edges):
    global image
    h = masked_edges.shape[0]/20
    # 霍夫变换检测车道线
    lines = cv2.HoughLinesP(masked_edges, 1, np.pi / 180, threshold=50, minLineLength=h, maxLineGap=h/4)

    # 绘制检测到的车道线
    line_image = np.zeros_like(image)
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            k = calculate_slope(line)
            if 10 >= k >= 0.5 or -10 <= k <= -0.5:  # 过滤斜率不正常的直线（接近水平或垂直）
                cv2.line(line_image, (x1, y1), (x2, y2), (0, 255, 0), 5)

    # 合并原图和车道线
    res_image = cv2.addWeighted(image, 0.8, line_image, 1, 0)
    return res_image


if __name__ == '__main__':
    image_path = "dataset/image5.jpg"
    image = cv2.imread(image_path)

    # 转换为灰度图像
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 将亮度过低的区域转换为黑色
    mask = gray_image < 150
    gray_image[mask] = 0

    # 应用高斯模糊
    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 3)

    cv2.imshow('origin', image)
    cv2.imshow('gray', gray_image)
    cv2.imshow('blurred', blurred_image)

    cv2.waitKey(0)  # 按任意键退出
    cv2.destroyAllWindows()


    # Canny 边缘检测
    edges = cv2.Canny(gray_image, 180, 255)
    edges_blurred = cv2.Canny(blurred_image, 180, 255)

    cv2.imshow('edge_origin', edges)
    cv2.imshow('edge_blurred', edges_blurred)

    cv2.waitKey(0)  # 按任意键退出
    cv2.destroyAllWindows()


    # 只保留ROI区域
    masked_edges = region_of_interest(edges_blurred)
    cv2.imshow('masked_edge',masked_edges)

    cv2.waitKey(0)  # 按任意键退出
    cv2.destroyAllWindows()


    # 显示结果
    final_image = detect_roadline(masked_edges)
    cv2.imshow('Lane Detection', final_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()